#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/2/27 13:56
# @Author  : 王凯
# @File    : tonghuashun_tools.py
# @Project : scrapy_spider

import psycopg2
from dbutils.pooled_db import PooledDB

from utils.db.mysqldb import MysqlDB
from utils.logs import logger

TONGHUASHUN_CONFIG = {
    "ip": "server-forward.local.wfq2020.com",
    "db": "tonghuashun-datasync",
    "port": 55432,
    "user_name": "z.developer.crawler.王凯",
    "user_pass": "2f2f1e791b70b53047af83fa289f87b0",
}


class PostgresqlDB(MysqlDB):
    def __init__(
        self,
        ip=TONGHUASHUN_CONFIG["ip"],
        port=TONGHUASHUN_CONFIG["port"],
        db=TONGHUASHUN_CONFIG["db"],
        user_name=TONGHUASHUN_CONFIG["user_name"],
        user_pass=TONGHUASHUN_CONFIG["user_pass"],
        **kwargs,
    ):
        super(PostgresqlDB).__init__(**kwargs)
        try:
            self.connect_pool = PooledDB(
                creator=psycopg2,
                mincached=1,
                maxcached=100,
                maxconnections=100,
                blocking=True,
                ping=7,
                host=ip,
                port=port,
                user=user_name,
                password=user_pass,
                database=db,
            )

        except Exception as e:
            logger.error(
                """
            连接失败：
            ip: {}
            port: {}
            db: {}
            user_name: {}
            user_pass: {}
            exception: {}
            """.format(
                    ip, port, db, user_name, user_pass, e
                )
            )
        else:
            logger.debug("连接到postgresql数据库 %s : %s" % (ip, db))

    def get_company_name_or_taxpayer_id(self, company_name_list):
        if isinstance(company_name_list, str):
            company_name_list = [company_name_list]

        result = {}
        taxpayer_id_list = []
        batch_size = 10

        for i in range(0, len(company_name_list), batch_size):
            coms = company_name_list[i: (i + batch_size)]
            taxpayer_id_list.extend(self._sub_get_company_name_or_taxpayer_id(coms))

        for k in company_name_list:
            tmp_res = [i for i in taxpayer_id_list if i["taxpayer_id"] == k or i["company_name"] == k]
            tmp_item = tmp_res[0] if tmp_res else {}
            result[k] = (
                {k: (v or "").replace("（", "(").replace("）", ")") for k, v in tmp_item.items()} if tmp_item else {}
            )

        return result

    def _sub_get_company_name_or_taxpayer_id(self, coms):
        fullname = ",".join([f"'{i}'" for i in coms])
        if fullname:
            sql = """ SELECT corp_name AS company_name, unified_social_credit_code AS taxpayer_id FROM enterprise_basic_info  WHERE (corp_name IN ({com_name_list_str}) OR unified_social_credit_code IN ({com_name_list_str})) AND isvalid = 1  UNION ALL  SELECT corp_name AS company_name, unified_social_credit_code AS taxpayer_id FROM individual_business_info  WHERE (corp_name IN ({com_name_list_str}) OR unified_social_credit_code IN ({com_name_list_str})) AND isvalid = 1""".format(
                com_name_list_str=fullname
            )
            taxpayer_id_list = self.find(sql, to_json=True)
        else:
            taxpayer_id_list = []
        return taxpayer_id_list


if __name__ == "__main__":
    test = PostgresqlDB()
    print(
        test.get_company_name_or_taxpayer_id(
            [
                "东台市孙之祥建材经营部",
                "怀仁市怀德商贸有限公司",
                "昆山市巴城镇葛裕成蟹舫",
                "湖南中涂化工有限责任公司",
                "丽水市易安机电工程有限公司",
                "北京凯远兴商贸公司",
                "献县霞光教育培训学校有限公司",
                "陕西正雅建筑设计有限公司",
                "南京艺森品至软装设计有限公司",
                "营口迪尼维尔汽保设备有限公司",
                "蚌埠市蚌山区家贵酒店朝阳路店",
                "济宁义德商贸有限公司",
            ]
        )
    )
